//
// Created by Mechrevo on 2023/4/23.
//
#if NCNN_VULKAN
#include "ncnn/gpu.h"
bool is_gpu =  ncnn::get_gpu_count()>0;
#endif

#include "opencv2/opencv.hpp"
#include "ncnn/net.h"
#include <cmath>
#include "retinaface.h"
#include "mnet.25.id.h"
#include "mnet.25.mem.h"

using namespace mnet_25_opt_param_id;
using namespace std;
using namespace cv;
using namespace retina;


vector<int> retina::resize_letter(Mat &image, float new_shape) {
    int h = static_cast<int>(image.rows);
    int w = static_cast<int>(image.cols);
    float new_h = new_shape / h;
    float new_w = new_shape / w;
    float r = (new_h < new_w ? new_h : new_w);

    return {static_cast<int>(round(h * r)), static_cast<int>(round(w * r))};
    // return {static_cast<int>(round(h)), static_cast<int>(round(w))};
};

//static inline float intersection_area(const FaceObject &a, const FaceObject &b) {
//    cv::Rect_<float> inter = a.rect & b.rect;
//    return inter.area();
//}

static ncnn::Mat generate_anchors(int base_size, const ncnn::Mat &ratios, const ncnn::Mat &scales) {

    int num_ratios = ratios.w;
    int num_scales = scales.w;
    ncnn::Mat anchors;
    anchors.create(4, num_ratios * num_scales);

    const float cx = static_cast<float>(base_size) * 0.5f;
    const float cy = static_cast<float>(base_size) * 0.5f;

    for (int i = 0; i < num_ratios; i++) {

        float ar = ratios[i];

        int r_w = static_cast<int>(round(static_cast<float>(base_size) / sqrt(ar)));
        int r_h = static_cast<int>(round(static_cast<float>(r_w) * ar));
        for (int j = 0; j < num_scales; j++) {
            float scale = scales[j];
            float rs_w = (float) r_w * scale;

            float rs_h = (float) r_h * scale;
            float *anchor = anchors.row(i * num_scales + j);

            anchor[0] = cx - rs_w * 0.5f;
            anchor[1] = cy - rs_h * 0.5f;
            anchor[2] = cx + rs_w * 0.5f;
            anchor[3] = cy + rs_h * 0.5f;


        }
    }
    return anchors;
}

// 通过分桶的方式进行快速排序
static void qsort_descent_inplace(vector<FaceObject> &face_object, int left, int right) {

    int i = left;
    int j = right;

    float p = face_object[(left + right) / 2].prob;

    while (i <= j) {
        while (face_object[i].prob > p) {
            i++;
        }
        while (face_object[j].prob < p) {
            j--;
        }
        if (i <= j) {
            // 交换位置
            std::swap(face_object[i], face_object[j]);
            i++;
            j--;
        }
    }
#pragma omp parallel sections // 编译器优化项
    {
#pragma omp section
        {
            if (left < j) qsort_descent_inplace(face_object, left, j);
        }
#pragma omp section
        {
            if (i < right) qsort_descent_inplace(face_object, i, right);
        }
    }

}

static void nms_sorted_bbox(vector<FaceObject> &face_object, vector<int> &picked, float nms_threshold) {
    picked.clear(); // 清空picked vector

    const int num_face_obj = (int) face_object.size();

    vector<float> areas(num_face_obj);

    for (int i = 0; i < num_face_obj; i++) {
        areas[i] = face_object[i].rect.area(); // 取出每个框的面积
    }

    for (int i = 0; i < num_face_obj; i++) {
        const FaceObject &current_face = face_object[i];

        int keep = 1;
        for (const auto &obj: picked) { // 迭代picked内的数据
            const FaceObject &keep_object = face_object[obj];

            // 计算交幷比

            float inter_area;
            inter_area = retina::intersection_area(current_face, keep_object);
            float union_area = areas[i] + areas[obj] - inter_area;
            if (inter_area / union_area > nms_threshold) {
                keep = 0;
            }
        }
        if (keep) {
            picked.push_back(i);
        }
    }


}

static void qsort_descent_inplace(vector<FaceObject> &face_object) { //
    if (face_object.empty()) {
        return;
    }
    qsort_descent_inplace(face_object, 0, (int) face_object.size() - 1); // 重载快排

}

static void
generate_proposals(const ncnn::Mat &anchors, int feat_stride, const ncnn::Mat &score_blob, const ncnn::Mat &bbox_blob,
                   const ncnn::Mat &landmark_blob, float prob_threshold, std::vector<FaceObject> &face_object) {
    int w = score_blob.w;
    int h = score_blob.h;
    // generate face proposal from bbox deltas and shifted anchors
    const int num_anchors = anchors.h;

    for (int i = 0; i < num_anchors; i++) {
        const float *anchor = anchors.row(i);
        const ncnn::Mat score = score_blob.channel(i + num_anchors);
        const ncnn::Mat bbox = bbox_blob.channel_range(i * 4, 4);
        const ncnn::Mat landmark = landmark_blob.channel_range(i * 10, 10);

        // shifted anchor
        float anchor_y = anchor[1];

        float anchor_w = anchor[2] - anchor[0];
        float anchor_h = anchor[3] - anchor[1];

        for (int q = 0; q < h; q++) {
            float anchor_x = anchor[0];
            for (int j = 0; j < w; j++) {

                int index = q * w + j;

                float prob = score[index];

                if (prob >= prob_threshold) {
                    // apply center size
                    float dx = bbox.channel(0)[index];
                    float dy = bbox.channel(1)[index];
                    float dw = bbox.channel(2)[index];
                    float dh = bbox.channel(3)[index];

                    float cx = anchor_x + anchor_w * 0.5f;
                    float cy = anchor_y + anchor_h * 0.5f;

                    float pb_cx = cx + anchor_w * dx;
                    float pb_cy = cy + anchor_h * dy;

                    float pb_w = anchor_w * exp(dw);
                    float pb_h = anchor_h * exp(dh);

                    float x0 = pb_cx - pb_w * 0.5f;
                    float y0 = pb_cy - pb_h * 0.5f;
                    float x1 = pb_cx + pb_w * 0.5f;
                    float y1 = pb_cy + pb_h * 0.5f;

                    FaceObject obj;
                    obj.rect.x = x0;
                    obj.rect.y = y0;
                    obj.rect.width = x1 - x0 + 1;
                    obj.rect.height = y1 - y0 + 1;
                    obj.landmark[0].x = cx + (anchor_w + 1) * landmark.channel(0)[index];
                    obj.landmark[0].y = cy + (anchor_h + 1) * landmark.channel(1)[index];
                    obj.landmark[1].x = cx + (anchor_w + 1) * landmark.channel(2)[index];
                    obj.landmark[1].y = cy + (anchor_h + 1) * landmark.channel(3)[index];
                    obj.landmark[2].x = cx + (anchor_w + 1) * landmark.channel(4)[index];
                    obj.landmark[2].y = cy + (anchor_h + 1) * landmark.channel(5)[index];
                    obj.landmark[3].x = cx + (anchor_w + 1) * landmark.channel(6)[index];
                    obj.landmark[3].y = cy + (anchor_h + 1) * landmark.channel(7)[index];
                    obj.landmark[4].x = cx + (anchor_w + 1) * landmark.channel(8)[index];
                    obj.landmark[4].y = cy + (anchor_h + 1) * landmark.channel(9)[index];
                    obj.prob = prob;

                    face_object.push_back(obj);
                }
                anchor_x += (float) feat_stride;
            }
            anchor_y += (float) feat_stride;
        }

    }
}

cv::Mat retina::draw_face_objects(const cv::Mat &bgr, const std::vector<FaceObject> &faceobjects) {
    cv::Mat image = bgr.clone();

    for (size_t i = 0; i < faceobjects.size(); i++) {
        const FaceObject &obj = faceobjects[i];

        fprintf(stderr, "%.5f at %.2f %.2f %.2f x %.2f\n", obj.prob,
                obj.rect.x, obj.rect.y, obj.rect.width, obj.rect.height);

        cv::rectangle(image, obj.rect, cv::Scalar(0, 255, 0));
        cv::circle(image, obj.landmark[0], 2, cv::Scalar(0, 255, 255), -1);
        cv::circle(image, obj.landmark[1], 2, cv::Scalar(0, 255, 255), -1);
        cv::circle(image, obj.landmark[2], 2, cv::Scalar(0, 255, 255), -1);
        cv::circle(image, obj.landmark[3], 2, cv::Scalar(0, 255, 255), -1);
        cv::circle(image, obj.landmark[4], 2, cv::Scalar(0, 255, 255), -1);

        char text[256];
        sprintf(text, "%.1f%%", obj.prob * 100);

        int baseLine = 0;
        cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);

        int x = obj.rect.x;
        int y = obj.rect.y - label_size.height - baseLine;
        if (y < 0)
            y = 0;
        if (x + label_size.width > image.cols)
            x = image.cols - label_size.width;

        cv::rectangle(image, cv::Rect(cv::Point(x, y), cv::Size(label_size.width, label_size.height + baseLine)),
                      cv::Scalar(255, 255, 255), -1);

        cv::putText(image, text, cv::Point(x, y + label_size.height),
                    cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 0));
    }
    return image;
//    cv::Mat image_;
//    cv::resize(image, image_, cv::Size(0, 0), 0.5, 0.5);
//    cv::imshow("image", image_);
//    cv::waitKey(0);
}

float get_dist_p2l(cv::Point2f center, cv::Point2f point1, cv::Point2f point2) {
    /**
     * center 中心点  point1 点1  point2  点2
     * **/
    float A, B, C;
    A = point1.y - point2.y;
    B = point2.x - point1.x;
    C = point1.x * point2.y + point1.y * point2.x;
    float distance;
    distance = ((float) fabs(A * center.x + B * center.y + C)) / ((float) sqrtf(A * A + B * B));
    return distance;
}

int InspectFaceIncline(vector<FaceObject> &face_objects) {
    /**
     * 判断人脸斜视
     * 规则: 低头判断: 当鼻子距离嘴巴两个点比较近的时候就是低头
     *       抬头判断: 当鼻子点距离眼睛比较近的时候就是抬头
     *       左斜脸:  当鼻子点距离左嘴巴点和左眼睛点比较近的时候判断为左斜脸
     *       右斜脸: 当鼻子点距离又嘴巴点和有眼睛点比较近的时候判断为右斜脸
     * **/
    vector<int> right_eye_points = {0, 1};
    vector<int> left_eye_points = {2, 3};

    vector<int> nose_points = {4, 5};
    vector<int> right_mouth_points{6, 7};
    vector<int> left_mouth_points{8, 9};
    return 0;
    // auto landmark = face_objects[0].landmark;
    // 根据五个关键点来判断是否斜视
}

/*
 *  包含模型推理部分，包含模型后处理结果，不包含最终结果判定
 * */
void retina::retina_face_predict(cv::Mat &image, vector<FaceObject> &face_objects) {
    float image_area = image.cols * image.rows;

    // ncnn::Net* retina_face = getModel();

    auto start = std::chrono::high_resolution_clock::now();
    ncnn::Net retina_face;
    retina_face.opt.use_vulkan_compute = true;
    retina_face.load_param(mnet_25_opt_param_bin);
    retina_face.load_model(mnet_25_opt_bin);
    auto end = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);


    vector<int> new_size = resize_letter(image);
    int width = new_size[1];
    int height = new_size[0];
    float ratio = image.cols / width * 1.0f; // 计算缩放比例
    cv::Mat dstImage;
    cv::resize(image, dstImage, cv::Size(width, height), 0, 0, cv::INTER_LINEAR);



    // float ratio = new_size[2];
//    int width = image.cols;
//    int height = image.rows;
    // ncnn::Mat in = ncnn::Mat::from_pixels(image.data, ncnn::Mat::PIXEL_BGR2RGB, width, height);
    ncnn::Mat in = ncnn::Mat::from_pixels(dstImage.data, ncnn::Mat::PIXEL_BGR2RGB, width, height);

    const float prob_threshold = 0.9f;
    const float nms_threshold = 0.4f;

    ncnn::Extractor ex = retina_face.create_extractor();
    ex.input(mnet_25_opt_param_id::BLOB_data, in);
    std::vector<FaceObject> face_proposals;
    //stride 32
    {
        ncnn::Mat score_blob, bbox_blob, landmark_blob;
        ex.extract(mnet_25_opt_param_id::BLOB_face_rpn_cls_prob_reshape_stride32, score_blob);
        ex.extract(mnet_25_opt_param_id::BLOB_face_rpn_bbox_pred_stride32, bbox_blob);
        ex.extract(mnet_25_opt_param_id::BLOB_face_rpn_landmark_pred_stride32, landmark_blob);
        const int base_size = 16;
        const int feat_stride = 32;
        ncnn::Mat ratios(1);
        ratios[0] = 1.f;
        ncnn::Mat scales(2);
        scales[0] = 32.f;
        scales[1] = 16.f;
        ncnn::Mat anchors = generate_anchors(base_size, ratios, scales);
        vector<FaceObject> face_objects32;
        generate_proposals(anchors, feat_stride, score_blob, bbox_blob, landmark_blob, prob_threshold, face_objects32);
        face_proposals.insert(face_proposals.end(), face_objects32.begin(), face_objects32.end());

    };
    // stride 16
    {
        ncnn::Mat score_blob, bbox_blob, landmark_blob;
        ex.extract(mnet_25_opt_param_id::BLOB_face_rpn_cls_prob_reshape_stride16, score_blob);
        ex.extract(mnet_25_opt_param_id::BLOB_face_rpn_bbox_pred_stride16, bbox_blob);
        ex.extract(mnet_25_opt_param_id::BLOB_face_rpn_landmark_pred_stride16, landmark_blob);
        const int base_size = 16;
        const int feat_stride = 16;
        ncnn::Mat ratios(1);
        ratios[0] = 1.f;
        ncnn::Mat scales(2);
        scales[0] = 8.f;
        scales[1] = 4.f;

        ncnn::Mat anchors = generate_anchors(base_size, ratios, scales);

        vector<FaceObject> face_objects16;
        generate_proposals(anchors, feat_stride, score_blob, bbox_blob, landmark_blob, prob_threshold, face_objects16);

        face_proposals.insert(face_proposals.end(), face_objects16.begin(), face_objects16.end());

    }
    // stride 8
    {
        ncnn::Mat score_blob, bbox_blob, landmark_blob;
        ex.extract(mnet_25_opt_param_id::BLOB_face_rpn_cls_prob_reshape_stride8, score_blob);
        ex.extract(mnet_25_opt_param_id::BLOB_face_rpn_bbox_pred_stride8, bbox_blob);
        ex.extract(mnet_25_opt_param_id::BLOB_face_rpn_landmark_pred_stride8, landmark_blob);

        const int base_size = 16;
        const int feat_stride = 8;
        ncnn::Mat ratios(1);
        ratios[0] = 1.f;
        ncnn::Mat scales(2);
        scales[0] = 2.f;
        scales[1] = 1.f;

        ncnn::Mat anchors = generate_anchors(base_size, ratios, scales);

        vector<FaceObject> face_objects8;
        generate_proposals(anchors, feat_stride, score_blob, bbox_blob, landmark_blob, prob_threshold, face_objects8);
        face_proposals.insert(face_proposals.end(), face_objects8.begin(), face_objects8.end());
    }
    // sort all proposals by score from highest to lowest
    qsort_descent_inplace(face_proposals);

    // nms with threshold
    vector<int> picked;
    nms_sorted_bbox(face_proposals, picked, nms_threshold);

    int keep_n = (int) picked.size();

    face_objects.resize(keep_n);

    for (int i = 0; i < keep_n; i++) {

        face_objects[i] = face_proposals[picked[i]]; // 将nms之后框进行存储
        // clip to image size  切换 xy wh to x1,y1 x2,y2
        float x0 = face_objects[i].rect.x;
        float y0 = face_objects[i].rect.y;
        float x1 = x0 + face_objects[i].rect.width;
        float y1 = y0 + face_objects[i].rect.height;

        x0 = std::max(std::min(x0, (float) width - 1), 0.f);
        y0 = std::max(std::min(y0, (float) height - 1), 0.f);
        x1 = std::max(std::min(x1, (float) width - 1), 0.f);
        y1 = std::max(std::min(y1, (float) height - 1), 0.f);

        // face_objects[i].rect.x = x0;
        // face_objects[i].rect.y = y0;
        // face_objects[i].rect.width = x1 - x0;
        // face_objects[i].rect.height = y1 - y0;

        face_objects[i].rect.x = x0 * ratio;
        face_objects[i].rect.y = y0 * ratio;
        face_objects[i].rect.width = (x1 - x0) * ratio;
        face_objects[i].rect.height = (y1 - y0) * ratio;
        face_objects[i].width = width;
        face_objects[i].height = height;
        face_objects[i].costTime = "model load cost time:" + std::to_string(duration.count()) + " ms";
    }
    vector<FaceObject> new_face_objects;
    for (const auto &obj: face_objects) {
        if ((obj.rect.area() / image_area) > 0.02) {
            new_face_objects.push_back(obj);
        }
    }
    face_objects = new_face_objects;
    new_face_objects.clear();
    // ai 智能核检 做任务完成时间(有数据情况下，第一版模型出来时间)
    // 其余任务的实现方式
}


